import base64
import logging
import hashlib
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

class AESCipher:
    """AES加密/解密工具类
    
    支持AES-ECB模式加密和解密，使用PKCS7填充。
    """
    
    def __init__(self, key):
        """初始化AES加密解密对象
        
        Args:
            key: 加密密钥，将被转为32位密钥
        """
        self.block_size = AES.block_size
        # 确保密钥长度为32字节(256位)
        self.key = self._ensure_key_length(key)
    
    def _ensure_key_length(self, key):
        """确保密钥长度为32字节，如果不足则使用SHA-256哈希扩展"""
        if isinstance(key, str):
            key = key.encode('utf-8')
            
        if len(key) == 32:  # 密钥长度正好是32字节
            return key
        elif len(key) < 32:  # 密钥长度小于32字节，使用SHA-256扩展
            return hashlib.sha256(key).digest()
        else:  # 密钥长度大于32字节，截断
            return key[:32]
    
    def encrypt(self, raw):
        """加密文本数据
        
        Args:
            raw: 待加密的原始文本（字符串）
            
        Returns:
            返回base64编码的密文字符串
        """
        if raw is None or raw == '':
            return None
        
        try:
            # 如果输入是字符串，转换为字节
            if isinstance(raw, str):
                raw = raw.encode('utf-8')
                
            # 创建AES-ECB模式的加密器
            cipher = AES.new(self.key, AES.MODE_ECB)
            
            # 填充数据并加密
            padded_data = pad(raw, self.block_size)
            encrypted_bytes = cipher.encrypt(padded_data)
            
            # Base64编码并返回字符串
            return base64.b64encode(encrypted_bytes).decode('utf-8')
            
        except Exception as e:
            logging.error(f"AES加密失败: {str(e)}")
            return None
    
    def decrypt(self, enc):
        """解密加密数据
        
        Args:
            enc: base64编码的加密数据
            
        Returns:
            返回解密后的字符串
        """
        if enc is None or enc == '':
            return None
            
        try:
            # 如果输入是字符串，转换为字节
            if isinstance(enc, str):
                enc = enc.encode('utf-8')
                
            # Base64解码
            enc = base64.b64decode(enc)
            
            # 创建AES-ECB模式的解密器
            cipher = AES.new(self.key, AES.MODE_ECB)
            
            # 解密并移除填充
            decrypted_padded = cipher.decrypt(enc)
            decrypted = unpad(decrypted_padded, self.block_size)
            
            # 将解密后的字节转换为字符串
            return decrypted.decode('utf-8')
            
        except Exception as e:
            logging.error(f"AES解密失败: {str(e)}")
            return None 